

import os
from PIL import Image, ImageOps
import numpy as np
from tqdm import tqdm

img_folder = r"G:\SGG\segmmetation\jpg_257"
mask_folder = r"G:\SGG\segmmetation\mask_257"
ouput_img_folder = r"G:\SGG\segmmetation\data_enforce\image"
ouput_mask_folder = r"G:\SGG\segmmetation\data_enforce\imagePNG"

os.makedirs(ouput_img_folder, exist_ok=True)
os.makedirs(ouput_mask_folder, exist_ok=True)

rotated_and_h_mirror = True
rotated_and_v_mirror = True
mirror = True
for img in tqdm(os.listdir(img_folder), desc="Process img files", unit='files'):
    img_path = os.path.join(img_folder, img)
    img_name = os.path.basename(img_path).split('.')[0]

    mask_path = os.path.join(mask_folder, img).replace('jpg', 'png')
    image = Image.open(img_path)
    mask = Image.open(mask_path)

    # ensure the img size
    assert image.size == (512, 512)
    assert mask.size == (512, 512)

    if mirror:
        image_mirrored = ImageOps.mirror(image)
        mask_mirrored = ImageOps.mirror(mask)

        save_img_name = f'{img_name}_mirror.jpg'
        save_mask_name = f'{img_name}_mirror.png'
        save_img_path = os.path.join(ouput_img_folder, save_img_name)
        save_mask_path = os.path.join(ouput_mask_folder, save_mask_name)
        image_mirrored.save(save_img_path)
        mask_mirrored.save(save_mask_path)

    angles = [90, 180, 270]
    for angle in angles:
        image_rotated_angle = image.rotate(angle)
        mask_rotated_angle = mask.rotate(angle)

        save_img_name = f'{img_name}_{angle}.jpg'
        save_mask_name = f'{img_name}_{angle}.png'
        save_img_path = os.path.join(ouput_img_folder, save_img_name)
        save_mask_path = os.path.join(ouput_mask_folder, save_mask_name)
        image_rotated_angle.save(save_img_path)
        mask_rotated_angle.save(save_mask_path)

        if rotated_and_h_mirror:
            image_rotated_angle_mirrored = ImageOps.mirror(image_rotated_angle)
            mask_rotated_angle_mirrored = ImageOps.mirror(mask_rotated_angle)

            save_img_name = f'{img_name}_{angle}_h_mirror.jpg'
            save_mask_name = f'{img_name}_{angle}_h_mirror.png'
            save_img_path = os.path.join(ouput_img_folder, save_img_name)
            save_mask_path = os.path.join(ouput_mask_folder, save_mask_name)
            image_rotated_angle_mirrored.save(save_img_path)
            mask_rotated_angle_mirrored.save(save_mask_path)

        if rotated_and_v_mirror:
            image_rotated_angle_mirrored = ImageOps.flip(image_rotated_angle)
            mask_rotated_angle_mirrored = ImageOps.flip(mask_rotated_angle)

            save_img_name = f'{img_name}_{angle}_v_mirror.jpg'
            save_mask_name = f'{img_name}_{angle}_v_mirror.png'
            save_img_path = os.path.join(ouput_img_folder, save_img_name)
            save_mask_path = os.path.join(ouput_mask_folder, save_mask_name)
            image_rotated_angle_mirrored.save(save_img_path)
            mask_rotated_angle_mirrored.save(save_mask_path)
